import os
import warnings
from PIL import Image

# 禁用 DecompressionBombWarning 警告
warnings.simplefilter("ignore", Image.DecompressionBombWarning)

# 增加最大像素限制，防止 DecompressionBombError 错误
Image.MAX_IMAGE_PIXELS = 200000000  # 设置更高的像素限制（比如200M像素）


def resize_image(image_path, grid_size, output_folder):
    """根据拼接图的 grid_size 将图片缩小到拼接前的小图像大小"""
    # 打开拼接图
    img = Image.open(image_path)
    img_width, img_height = img.size

    # 计算每个小图像的原始大小
    small_img_width = img_width // grid_size
    small_img_height = img_height // grid_size

    # 缩小拼接图回到每个小图像的大小
    resized_img = img.resize((small_img_width, small_img_height))

    # 获取原始文件名，并生成保存路径
    image_name = os.path.basename(image_path)
    save_path = os.path.join(output_folder, image_name)

    # 保存结果图像
    resized_img.save(save_path)
    print(f"✅ 图像已保存: {save_path}")


def process_images(image_folder, concat_sizes, grid_sizes, output_folder):
    """遍历文件夹中的图像并按指定的拼接大小调整图像"""
    os.makedirs(output_folder, exist_ok=True)

    # 获取所有图像文件并提取出数字部分进行排序
    image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(('jpg', 'jpeg', 'png', 'webp'))]
    image_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))  # 按数字部分排序

    cumulative_images = 0
    grid_count = 0

    # 遍历每个拼接图像
    for image_file in image_files:
        image_path = os.path.join(image_folder, image_file)
        image_number = int(image_file.split('_')[1].split('.')[0])  # 提取图片编号（例如 0, 1, 2,...）
        # 如果图片编号小于当前累积数量，则使用当前的拼接方式（grid_size）
        if cumulative_images <= image_number < cumulative_images + concat_sizes[grid_count]:
            grid_size = grid_sizes[grid_count]
            resize_image(image_path, grid_size, output_folder)
        else:
            # 更新已处理图像数量和拼接方式索引
            cumulative_images += concat_sizes[grid_count]
            grid_count += 1
            # 处理当前拼接方式下的图像
            grid_size = grid_sizes[grid_count]
            resize_image(image_path, grid_size, output_folder)


if __name__ == "__main__":
    image_folder = r"F:\yolo_summary\yolo_code\car\521\test_merge\images"  # 输入图像文件夹
    output_folder = r"F:\yolo_summary\yolo_code\car\521\test_merge\images_resized"  # 输出文件夹

    # 第一个数组：拼接的张数
    concat_sizes = [20, 12, 7]

    grid_sizes = [3, 4, 5]  # 例如：2x2，3x3，4x4 等拼接方式

    # 调用处理函数
    process_images(image_folder, concat_sizes, grid_sizes, output_folder)
